from .layers import *

import os
import numpy as np
import tensorflow as tf

initializer = tf.contrib.layers.xavier_initializer()


def miccai2018_brat(inputs, training, keep, num_layers=4):
    # https://zhuanlan.zhihu.com/p/71578701
    base = 32  # 32 on Titan V100 32G, 16 on Titan V100 16G
    encoder_features = [base * 2 ** i for i in range(num_layers + 1)]
    features = encoder_features + encoder_features[::-1]
    features = [[f, f] for f in features]

    shape = inputs.shape.as_list()
    shape[0] = 1
    inputs = tf.reshape(inputs, shape)

    concat_tensor = {}
    tensor = Conv(inputs, base, 3, 1, keep, 'cd')

    for layer in range(0, num_layers):
        feature = features[1 + layer]
        concat_tensor[layer] = tensor
        tensor = MaxPool3dConv(tensor, feature[0])
        tensor = preActive(tensor, keep, feature)

    tensor = preActive(tensor, keep, features[1 + num_layers])

    for layer in range(num_layers - 1, -1, -1):
        #upsampled = Upsamping3d(tensor)
        upsampled = NNUpsamping3d(tensor, keep)
        concated = tf.concat([concat_tensor[layer], upsampled], -1)
        feature = features[1 + num_layers * 2 - layer]
        tensor = preActive(concated, keep, feature)

    last_conv = tf.layers.conv3d(tensor, 1, 1, padding="same",
                                 use_bias=True, activation=None,
                                 kernel_initializer=initializer)
    return last_conv
